from crewai.flow.flow import Flow, start, listen, router

class RouterFlow(Flow):

    @start()
    def start_node(self):
        self.state["msg"] = "你好，这是开始节点"

    @listen(start_node)
    def step_one(self):
        self.state["step_one"] = "这是第一个节点"

    @router(step_one)
    def router_node(self):
        print(f"路由节点收到: {self.state}")
        if True:
            return 'a'
        else:
            return 'b'
    @listen('a')
    def branch_one(self):
        self.state["branch"] = "这是路由分支一"
    @listen('b')
    def branch_two(self):
        self.state["branch"] = "这是路由分支二"



flow = RouterFlow()
flow.plot("router_flow.png")
flow.kickoff()
print("流程最终输出:", flow.state)
